# Get Noto font
# 注意：fonts-noto-cjk 不是 Python 包，无需通过 pip 安装，需先在系统中安装字体
import requests
from PIL import Image, ImageDraw, ImageFont
from io import BytesIO
from bs4 import BeautifulSoup, Tag
import re

# Function to draw bounding boxes and text on images based on HTML content
def draw_bbox_html(image_path, full_predict):
    """
    可视化 Qwenvl HTML 的 data-bbox 框并展示文本，坐标为相对 0-1000。
    过滤规则：跳过 <ol>，仅绘制 <li> 子项和其它元素。
    """
    # 读取图片
    if image_path.startswith("http"):
        response = requests.get(image_path)
        image = Image.open(BytesIO(response.content)).convert("RGB")
    else:
        image = Image.open(image_path).convert("RGB")
    width = image.width
    height = image.height

    soup = BeautifulSoup(full_predict, 'html.parser')
    elements_with_bbox = soup.find_all(attrs={'data-bbox': True})

    # 保留原过滤逻辑
    filtered_elements = []
    for el in elements_with_bbox:
        if el.name == 'ol':
            continue  # 跳过 <ol>
        elif el.name == 'li' and el.parent.name == 'ol':
            filtered_elements.append(el)  # 仅保留 <ol> 下的 <li>
        else:
            filtered_elements.append(el)

    # 字体兼容：优先加载 Noto Serif CJK SC，失败则降级
    font = None
    try:
        # 尝试加载 Noto Serif CJK SC（需确保系统已安装该字体）
        font = ImageFont.truetype("Noto Serif CJK SC", 10)
    except Exception:
        try:
            # 备选：加载 Noto Sans CJK 系列
            font = ImageFont.truetype("NotoSansCJK-Regular.ttc", 10)
        except Exception:
            # 最终降级为默认字体
            font = ImageFont.load_default()
    draw = ImageDraw.Draw(image)
    
    # 绘制框与文本
    for element in filtered_elements:
        bbox_str = element['data-bbox']
        text = element.get_text(strip=True)
        try:
            x1, y1, x2, y2 = map(int, bbox_str.split())
        except Exception:
            continue

        bx1 = int(x1 / 1000 * width)
        by1 = int(y1 / 1000 * height)
        bx2 = int(x2 / 1000 * width)
        by2 = int(y2 / 1000 * height)
        
        if bx1 > bx2:
            bx1, bx2 = bx2, bx1
        if by1 > by2:
            by1, by2 = by2, by1
            
        draw.rectangle([bx1, by1, bx2, by2], outline='red', width=2)
        draw.text((bx1, by2), text, fill='black', font=font)

    image.show()

    
# Function to draw bounding boxes on images based on Markdown content
def draw_bbox_markdown(image_path, md_content):
    """
    只可视化Markdown中的 <!-- Image/Table (x1, y1, x2, y2) --> 坐标框，坐标为相对0-1000
    Table 用绿色框，Image 用蓝色框。
    """
    if image_path.startswith("http"):
        response = requests.get(image_path)
        image = Image.open(BytesIO(response.content)).convert("RGB")
    else:
        image = Image.open(image_path).convert("RGB")
    width = image.width
    height = image.height

    pattern = r"<!-- (Image|Table) \(\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\) -->"
    matches = re.findall(pattern, md_content)
    draw = ImageDraw.Draw(image)
    for item in matches:
        typ, x1, y1, x2, y2 = item
        x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
        bx1 = int(x1 / 1000 * width)
        by1 = int(y1 / 1000 * height)
        bx2 = int(x2 / 1000 * width)
        by2 = int(y2 / 1000 * height)
        if bx1 > bx2:
            bx1, bx2 = bx2, bx1
        if by1 > by2:
            by1, by2 = by2, by1
        color = 'blue' if typ == "Image" else 'red'
        draw.rectangle([bx1, by1, bx2, by2], outline=color, width=6)

    image.show()


# Function to clean and format HTML content
def clean_and_format_html(full_predict):
    soup = BeautifulSoup(full_predict, 'html.parser')
    
    # Regular expression pattern to match 'color' styles in style attributes
    color_pattern = re.compile(r'\bcolor:[^;]+;?')

    # Find all tags with style attributes and remove 'color' styles
    for tag in soup.find_all(style=True):
        original_style = tag.get('style', '')
        new_style = color_pattern.sub('', original_style)
        if not new_style.strip():
            del tag['style']
        else:
            new_style = new_style.rstrip(';')
            tag['style'] = new_style
            
    # Remove 'data-bbox' and 'data-polygon' attributes from all tags
    for attr in ["data-bbox", "data-polygon"]:
        for tag in soup.find_all(attrs={attr: True}):
            del tag[attr]

    classes_to_update = ['formula.machine_printed', 'formula.handwritten']
    # Update specific class names in div tags
    for tag in soup.find_all(class_=True):
        if isinstance(tag, Tag) and 'class' in tag.attrs:
            new_classes = [cls if cls not in classes_to_update else 'formula' for cls in tag.get('class', [])]
            tag['class'] = list(dict.fromkeys(new_classes))  # Deduplicate and update class names

    # Clear contents of divs with specific class names and rename their classes
    for div in soup.find_all('div', class_='image caption'):
        div.clear()
        div['class'] = ['image']

    classes_to_clean = ['music sheet', 'chemical formula', 'chart']
    # Clear contents and remove 'format' attributes of tags with specific class names
    for class_name in classes_to_clean:
        for tag in soup.find_all(class_=class_name):
            if isinstance(tag, Tag):
                tag.clear()
                if 'format' in tag.attrs:
                    del tag[attr]

    # Manually build the output string
    output = []
    for child in soup.body.children:
        if isinstance(child, Tag):
            output.append(str(child))
            output.append('\n')  # Add newline after each top-level element
        elif isinstance(child, str) and not child.strip():
            continue  # Ignore whitespace text nodes
    complete_html = f"""```html\n<html><body>\n{" ".join(output)}</body></html>\n```"""
    return complete_html